import jieba
import re
import json

# 分词与文本处理工具，支持中文分词、数字转换、停用词过滤等功能

# 定义停用词表
stopwords = set(['的', '是', '啊'])

# 自定义词典模板
json_template = {"intent": "", "target": -1}

# 中文数字转阿拉伯数字的字典
chinese_to_arabic = {
    '零': 0, '一': 1, '二': 2, '三': 3, '四': 4, '五': 5,
    '六': 6, '七': 7, '八': 8, '九': 9, '十': 10,
    '百': 100, '千': 1000, '万': 10000, '亿': 100000000
}

# 水果中英文对照字典（本例未用到，可忽略）
fruit_name_mapping = {
    '苹果': 'apple', '香蕉': 'banana', '橙子': 'orange',
    '橘子': 'orange', '草莓': 'strawberry', '西瓜': 'watermelon', '菠萝': 'pineapple'
}

# 通过字典从汉字中提取数字
def chinese_to_arabic_number(chinese_num):
    # 从后往前遍历中文数字字符串，累加计算对应的阿拉伯数字
    total = 0
    r = 1
    for i in range(len(chinese_num) - 1, -1, -1):
        val = chinese_to_arabic[chinese_num[i]]
        if val == 10000 or val == 100000000:
            if r < val:
                r *= val
            else:
                r //= val
        else:
            total += r * val
    return total

# 将代表数字的中文汉字替换成阿拉伯数字
def replace_chinese_numbers(text):
    pattern = re.compile(r'零?[一二三四五六七八九十百千万亿]+')
    matches = pattern.findall(text)
    for match in matches:
        arabic_number = chinese_to_arabic_number(match)
        text = text.replace(match, str(arabic_number))
    return text

# 命令字符串预处理（去杂、转英文、分词、去停用词）
def preprocess(text):
    text = re.sub(r'\s+', '', text)  # 去除空白字符
    text = replace_chinese_numbers(text)  # 中文数字转阿拉伯数字
    text = text.lower()  # 转小写
    tokens = list(jieba.lcut(text))  # 分词
    filtered_tokens = [token for token in tokens if token not in stopwords]  # 去停用词
    # 替换中文水果名（本例未用到，可忽略）
    for i, token in enumerate(filtered_tokens):
        if token in fruit_name_mapping:
            filtered_tokens[i] = fruit_name_mapping[token]
    return filtered_tokens

# 从词元中提取意图和目标编号
def extract_intent(tokens):
    intents = {
        'go_to': ['移动', '去', '前往'],
        'pick': ['抓取', '拿起', '抓', '拿', '取'],  # "拿起"对应pick意图
        'release': ['放下', '放到', '放', '放置']  # "放到"对应release意图
    }
    result = []
    i = 0
    while i < len(tokens):
        for intent, triggers in intents.items():
            if tokens[i] in triggers:
                # 寻找目标编号（第一个遇到的数字）
                j = i + 1
                while j < len(tokens) and not (tokens[j].isdigit() or tokens[j] == '号'):
                    j += 1
                target_index = -1
                if j < len(tokens):
                    for k in range(j, len(tokens)):
                        if tokens[k].isdigit():
                            target_index = int(tokens[k])
                            break
                # 生成JSON结果
                json_obj = json_template.copy()
                json_obj['intent'] = intent
                json_obj['target'] = target_index
                result.append(json_obj)
                i = j  # 跳过已处理的词元
                break
        else:
            i += 1
    return result

# 主函数：从原始文本中提取意图
def get_intent_from_text(text):
    filtered_input = preprocess(text)
    return extract_intent(filtered_input)

# -------------------- 测试示例 --------------------
if __name__ == "__main__":
    # 待处理的复合命令（包含"拿起"和"放下"）
    command = "拿起1号放到7号抓取2号放到8号抓取3号放到9号抓取4号放到10号"
    # 调用主函数处理
    result = get_intent_from_text(command)
    # 输出结构化结果（JSON格式）
    print("处理结果：")
    for item in result:
        print(json.dumps(item, ensure_ascii=False))


